
//crypto.c文件
#include <stdio.h>
#include <string.h>
#include <mbedtls/aes.h>
#include <mbedtls/arc4.h>
#include <mbedtls/md5.h>
#include <mbedtls/sha1.h>
#include <mbedtls/sha256.h>
#include <mbedtls/sha512.h>
#include "crypto.h"

static int __aes_cbc_encrypt(uint8_t *data, int len,
        const uint8_t *key, const uint8_t *iv,
        int key_size, int pad_type)
{
    int i, length;
    uint8_t ivcpy[16];
    uint8_t pad = 0;
    mbedtls_aes_context ctx;

    length = (len + 0x0f) & (~0x0f);

    if (pad_type == CRYPTO_AES_PADDING_NOPADDING && (len & 0x0f))
        return -1;

    if ((len & 0x0f) == 0 && pad_type == CRYPTO_AES_PADDING_PKCS5PADDING)
        length += 16;

    if (pad_type == CRYPTO_AES_PADDING_PKCS5PADDING)
        pad = length - len;

    for (i=0; i<length - len; i++)
        data[len + i] = pad;

    len = length;

    memcpy(ivcpy, iv, 16);

    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_enc(&ctx, key, key_size);
    while (length > 0) {
        for (i=0; i<16; i++)
            data[i] ^= ivcpy[i];
        mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_ENCRYPT, data, data);
        memcpy(ivcpy, data, 16);
        data += 16; length -= 16;
    }
    mbedtls_aes_free(&ctx);

    return len;
}

static int __aes_cbc_decrypt(uint8_t *data, int len,
        const uint8_t *key, const uint8_t *iv,
        int key_size, int pad_type)
{
    int i, length;
    uint8_t ivcpy[16];
    uint8_t *ptr = data;
    uint8_t pad;
    mbedtls_aes_context ctx;

    if (len & 0xf)
        return -1;

    length = len;

    memcpy(ivcpy, iv, 16);

    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_dec(&ctx, key, key_size);

    while (length > 0) {
        uint8_t tmp[16];

        memcpy(tmp, data, 16);
        mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_DECRYPT, data, data);
        for (i=0; i<16; i++)
            data[i] ^= ivcpy[i];
        memcpy(ivcpy, tmp, 16);
        data += 16; length -= 16;
    }

    mbedtls_aes_free(&ctx);

    if (pad_type == CRYPTO_AES_PADDING_NOPADDING || \
            pad_type == CRYPTO_AES_PADDING_ZEROPADDING)
        return len;

    pad = ptr[len - 1];
    if (pad < 1 || pad > 16)
        return -1;

    for (i=1; i<=pad; i++) {
        if (ptr[len - i] != pad)
            return -1;
        ptr[len - i] = 0;
    }

    return len - pad;
}

int crypto_aes_cbc_encrypt(const uint8_t *input, int len,
        uint8_t *output, const uint8_t *key, const uint8_t *iv,
        int key_size, int pad_type)
{
    if (key_size != CRYPTO_AES_KEY_128BITS && \
            key_size != CRYPTO_AES_KEY_256BITS)
        return -1;

    if (pad_type != CRYPTO_AES_PADDING_NOPADDING && \
            pad_type != CRYPTO_AES_PADDING_ZEROPADDING && \
            pad_type != CRYPTO_AES_PADDING_PKCS5PADDING)
        return -1;

    memcpy(output, input, len);

    return __aes_cbc_encrypt(output, len, key, iv, key_size, pad_type);
}

int crypto_aes_cbc_decrypt(const uint8_t *input, int len,
        uint8_t *output, const uint8_t *key, const uint8_t *iv,
        int key_size, int pad_type)
{
    if (key_size != CRYPTO_AES_KEY_128BITS && \
            key_size != CRYPTO_AES_KEY_256BITS)
        return -1;

    if (pad_type != CRYPTO_AES_PADDING_NOPADDING && \
            pad_type != CRYPTO_AES_PADDING_ZEROPADDING && \
            pad_type != CRYPTO_AES_PADDING_PKCS5PADDING)
        return -1;

    memcpy(output, input, len);

    return __aes_cbc_decrypt(output, len, key, iv, key_size, pad_type);
}

int crypto_aes_ecb_encrypt(const uint8_t *input, int len,
        uint8_t *output, const uint8_t *key, int keybits)
{
    int i, cnt;
    mbedtls_aes_context ctx;

    if (len % 16)
        return -1;

    cnt = len / 16;

    if (keybits != CRYPTO_AES_KEY_128BITS && \
            keybits != CRYPTO_AES_KEY_256BITS)
        return -1;

    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_enc(&ctx, key, keybits);
    for (i=0; i<cnt; i++)
        mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_ENCRYPT, \
                input + (i << 4), output + (i << 4));

    return len;
}

int crypto_aes_ecb_decrypt(const uint8_t *input, int len,
        uint8_t *output, const uint8_t *key, int keybits)
{
    int i, cnt;
    mbedtls_aes_context ctx;

    if (len % 16)
        return -1;

    cnt = len / 16;

    if (keybits != CRYPTO_AES_KEY_128BITS && \
            keybits != CRYPTO_AES_KEY_256BITS)
        return -1;

    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_dec(&ctx, key, keybits);
    for (i=0; i<cnt; i++)
        mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_DECRYPT, \
                input + (i << 4), output + (i << 4));

    return len;
}

/*
static int protocol_data_encrypt(uint8_t *data,uint8_t *output,int len)//加密函数
{
    int encrypt;
    encrypt = (len + 0xf) & (~0xf);

    encrypt = crypto_aes_cbc_encrypt(data, len, output, pub_aes_key, initial_vector, \
            CRYPTO_AES_KEY_128BITS, CRYPTO_AES_PADDING_PKCS5PADDING);

    return encrypt;
}

static int protocol_data_decrypt(uint8_t *intput,uint8_t *output, int len)//解密函数
{
    int decrypt = len;

    if (len & 0x0f)
        return -1;

    decrypt = crypto_aes_cbc_decrypt(intput, len, output, pub_aes_key, initial_vector, \
            CRYPTO_AES_KEY_128BITS, CRYPTO_AES_PADDING_PKCS5PADDING);
    if (decrypt < 0)
        return -1;

    return decrypt;
}
*/



int crypto_rc4(const uint8_t *input, int len,
        uint8_t *output, const uint8_t *key, int key_size)
{
    mbedtls_arc4_context ctx;

    mbedtls_arc4_init(&ctx);
    mbedtls_arc4_setup(&ctx, key, key_size);
    mbedtls_arc4_crypt(&ctx, len, input, output);
    mbedtls_arc4_free(&ctx);

    return len;
}

int crypto_md5(uint8_t *data, int len, uint8_t dest[16])
{
    mbedtls_md5(data, len, dest);

    return 16;
}

int crypto_sha1(uint8_t *data, int len, uint8_t dest[20])
{
    mbedtls_sha1(data, len, dest);

    return 20;
}

int crypto_sha224(uint8_t *data, int len, uint8_t dest[28])
{
    mbedtls_sha256(data, len, dest, 1);

    return 28;
}

int crypto_sha256(uint8_t *data, int len, uint8_t dest[32])
{
    mbedtls_sha256(data, len, dest, 0);

    return 32;
}

int crypto_sha384(uint8_t *data, int len, uint8_t dest[48])
{
    mbedtls_sha512(data, len, dest, 1);

    return 48;
}

int crypto_sha512(uint8_t *data, int len, uint8_t dest[64])
{
    mbedtls_sha512(data, len, dest, 0);

    return 64;
}




